# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, ClassVar, Optional, Self, Type

from typing_extensions import override

from crossbench.parse import ObjectParser
from crossbench.probes.json import JsonResultProbe, JsonResultProbeContext
from crossbench.probes.metric import MetricsMerger

if TYPE_CHECKING:
  import datetime as dt

  from crossbench import exception
  from crossbench.probes.probe import ProbeConfigParser, ProbeKeyT
  from crossbench.probes.results import ProbeResult
  from crossbench.runner.actions import Actions
  from crossbench.runner.groups.browsers import BrowsersRunGroup
  from crossbench.runner.groups.stories import StoriesRunGroup
  from crossbench.types import Json


def parse_javascript(value: str) -> str:
  # TODO: maybe add more sanity checks
  return ObjectParser.non_empty_str(value, name="javascript")


class JSProbe(JsonResultProbe):
  """
  Probe for extracting arbitrary metrics using custom javascript code.
  """
  NAME: ClassVar = "js"

  @classmethod
  @override
  def config_parser(cls) -> ProbeConfigParser[Self]:
    parser = super().config_parser()
    parser.add_argument(
        "setup",
        type=parse_javascript,
        help=(
            "Optional JavaScript code that is run immediately before a story. "
            "This can be used for setting up some JS tracking code or patch "
            "existing code for custom metric tracking."))
    parser.add_default_argument(
        "js",
        type=parse_javascript,
        help=("Required JavaScript code that is run immediately after "
              "a story has finished. The code must return a JS object with "
              "(nested) metric values (numbers)."))
    return parser

  def __init__(self, js: Optional[str], setup: Optional[str] = None) -> None:
    super().__init__()
    self._setup_js = setup
    self._metric_js = js

  @property
  def setup_js(self) -> Optional[str]:
    return self._setup_js

  @property
  def metric_js(self) -> Optional[str]:
    return self._metric_js

  @property
  @override
  def key(self) -> ProbeKeyT:
    return super().key + (
        ("setup_js", self._setup_js),
        ("metric_js", self._metric_js),
    )

  def get_context_cls(self) -> Type[JSProbeContext]:
    return JSProbeContext

  def merge_stories(self, group: StoriesRunGroup) -> ProbeResult:
    merged = MetricsMerger.merge_json_list(
        story_group.results[self].json
        for story_group in group.repetitions_groups)
    return self.write_group_result(group, merged)

  def merge_browsers(self, group: BrowsersRunGroup) -> ProbeResult:
    return self.merge_browsers_json_list(group).merge(
        self.merge_browsers_csv_list(group))


class JSProbeContext(JsonResultProbeContext[JSProbe]):

  def __init__(self, *args, **kwargs) -> None:
    super().__init__(*args, **kwargs)
    self._json_data: dict[str, Any] = {}

  def _update_metrics(self, data: dict[str, Any]) -> None:
    existing_metrics = {k: v for k, v in self._json_data.items() if k in data}
    logging.debug("JSProbe: The following metrics will be overwritten: %s",
                  existing_metrics)
    self._json_data.update(data)

  @override
  def to_json(self, actions: Actions) -> Json:
    if metric_js := self.probe.metric_js:
      data = actions.js(metric_js)
      self._update_metrics(ObjectParser.non_empty_dict(data, "JS metric data"))
    return self._json_data

  def invoke(self, info_stack: exception.TInfoStack, timeout: dt.timedelta,
             **kwargs) -> None:
    del info_stack
    js = kwargs.pop("js")
    self.expect_no_extra_kwargs(kwargs)
    with self.run.actions(
        f"Probe({self.probe.name}) invoke", measure=False) as actions:
      data = actions.js(js, timeout=timeout)
      self._update_metrics(ObjectParser.non_empty_dict(data, "JS metric data"))

  def start(self) -> None:
    if setup_js := self.probe.setup_js:
      with self.run.actions(f"Probe({self.probe.name}) setup") as actions:
        actions.js(setup_js)
